import cv2
import pickle
import numpy as np
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt

# 读取原图和原分割图
original = cv2.imread("Sandstone_2.tif", cv2.IMREAD_GRAYSCALE)
segment = cv2.imread("Sandstone_2_segment.tif", cv2.IMREAD_GRAYSCALE)

# 加载保存的模型
with open('clf_gaussian.pkl', 'rb') as model_file:
    clf = pickle.load(model_file)

# 推理分割图
chosen_algorithm = 'gaussian'  # 选择与训练时一致的特征提取算法
if chosen_algorithm == 'mean':
    processed_image_test = cv2.blur(original, (5, 5))  # 5x5均值滤波
elif chosen_algorithm == 'gaussian':
    processed_image_test = cv2.GaussianBlur(original, (5, 5), 0)  # 高斯滤波
elif chosen_algorithm == 'sobel':
    processed_image_test = cv2.Sobel(original, cv2.CV_64F, 1, 1, ksize=5)  # Sobel边缘检测
elif chosen_algorithm == 'canny':
    processed_image_test = cv2.Canny(original, 50, 150)  # Canny边缘检测

# 将处理后的图像转换为一维数组
X_test = processed_image_test.flatten()

# 使用训练好的模型进行预测
y_pred = clf.predict(X_test.reshape(-1, 1))

# 将预测结果还原为图像形状
segmentation = y_pred.reshape(original.shape)

# 计算准确率
accuracy = accuracy_score(segment.flatten(), y_pred)

# 创建 Matplotlib subplot
fig, axs = plt.subplots(1, 3, figsize=(15, 5))

# 显示原始图像
axs[0].imshow(original, cmap='gray')
axs[0].set_title("Original image")

# 显示原始分割图像
axs[1].imshow(segment, cmap='gray')
axs[1].set_title("Segement")

# 显示推理分割图像
axs[2].imshow(segmentation, cmap='gray')
axs[2].set_title(f"Segmentation (acc: {accuracy:.3f})")


# 显示图像
plt.show()
